import os

import torch
import warnings
warnings.simplefilter("ignore") # Stop spam of future warnings I'm seeing
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def clip_vit_b32():
    clip_path = '/data/models/clip-ViT-B-32'
    from sentence_transformers import SentenceTransformer
    from PIL import Image

    #Load CLIP model
    model = SentenceTransformer(clip_path,device='cuda')
    return model

    #Encode an image:
    img_emb = model.encode(Image.open('two_dogs_in_snow.jpg'))

    #Encode text descriptions
    text_emb = model.encode(['Two dogs in the snow', 'A cat on a table', 'A picture of London at night'])

    #Compute cosine similarities 
    cos_scores = util.cos_sim(img_emb, text_emb)
    print(cos_scores)

def get_embedding(clip_model, text_or_imgpath):
    vector = clip_model.encode(text_or_imgpath)
    return vector


def test():
    import torch.nn.functional as F

    clip_model = clip_vit_b32()
    text = ('This garment is a beige blazer with a structured design. \
            It features a single-breasted front with two buttons, a notched lapel, \
                and a breast pocket. The blazer has a clean, minimalist aesthetic with a \
                smooth texture and subtle herringbone pattern. It is suitable for fall or \
                    winter seasons due to its neutral tone and layered look. The cuffs are buttoned, \
                        the neckline is a classic notch, and it has two flap pockets. The blazer is made of \
                            a smooth fabric, likely a blend of materials such as cotton or polyester, providing \
                                both comfort and durability. The overall design is elegant and versatile, making it \
                                    appropriate for both professional and casual settings.')
    text2 = 'gray clothing'
    txt_emb = get_embedding(clip_model, text)
    txt_emb2 = get_embedding(clip_model, text2)

    similarity = F.cosine_similarity(txt_emb, txt_emb2)
    print("Cosine similarity between text and image embeddings:", similarity.item())


if __name__ == '__main__':
    import argparse,os
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda


    test()
